package com.jboltai.resource.embedding;


import cn.hutool.core.util.StrUtil;
import com.jboltai.capability.embedding.EmbeddingProcessor;
import com.jboltai.capability.embedding.ali.TongyiEmbeddingImpl;
import com.jboltai.capability.embedding.baichuan.BaichuanEmbeddingImpl;
import com.jboltai.capability.embedding.baidu.WenxinEmbeddingImpl;
import com.jboltai.capability.embedding.local.LocalEmbeddingImpl;
import com.jboltai.capability.embedding.openai.OpenAIEmbeddingImpl;

import java.util.concurrent.ConcurrentHashMap;

/**
 * 向量化模型
 */
public class EmbeddingModel {

    public static final String NAME_BAI_CHUAN = "Baichuan-Text-Embedding";
    public static final String NAME_OPENAI_ADA002 = "text-embedding-ada-002";
    public static final String NAME_OPENAI_3SMALL = "text-embedding-3-small";

    public static final String NAME_TONGYI_V3 = "text-embedding-v3";
    public static final String NAME_TONGYI_V2 = "text-embedding-v2";

    public static final String NAME_BGE_LARGE_ZH_LOCAL = "bge-large-zh-local";

    public static final String NAME_OLLAMA = "ollama";
    /**
     * 硅基流动
     */
    public static final String NAME_SILICONFLOW = "siliconflow";

    public static final String NAME_BGE_LARGE_ZH = "bge-large-zh";
    public static final String NAME_BGE_LARGE_EN = "bge-large-en";

    public static final String NAME_BAIDU_QIANFAN = "embedding-v1";

    public static final String NAME_TENCENT_BGE_BASE_ZH = "Tencent-BGE_BASE_ZH";
    public static final String NAME_TENCENT_M3E_BASE = "Tencent-M3E_BASE";
    public static final String NAME_TENCENT_E5_LARGE_V2 = "Tencent-E5_LARGE_V2";
    public static final String NAME_TENCENT_TEXT2VEC_LARGE_CHINESE = "Tencent-TEXT2VEC_LARGE_CHINESE";
    public static final String NAME_TENCENT_MULTILINGUAL_E5_BASE = "Tencent-MULTILINGUAL_E5_BASE";


    /**
     * 百川 embedding模型
     */
    public static final EmbeddingModel BAI_CHUAN = new EmbeddingModel("百川Embedding", NAME_BAI_CHUAN, 1024, BaichuanEmbeddingImpl.INSTANCE, 500, 16);
    /**
     * 阿里通义 embedding模型
     */
    public static final EmbeddingModel TONGYI_V3 = new EmbeddingModel("通义V3Embedding", NAME_TONGYI_V3,
            1024, TongyiEmbeddingImpl.INSTANCE, 8192, 6);
    /**
     * 阿里通义 embedding模型
     */
    public static final EmbeddingModel TONGYI_V2 = new EmbeddingModel("通义V2Embedding", NAME_TONGYI_V2,
            1536, TongyiEmbeddingImpl.INSTANCE, 2048, 25);

    /**
     * openai embedding模型
     */
    public static final EmbeddingModel OPENAI_ADA002 = new EmbeddingModel("OpenAIEmbedding", NAME_OPENAI_ADA002, 1536, OpenAIEmbeddingImpl.INSTANCE, 500, 16);
    public static final EmbeddingModel OPENAI_3_SMALL = new EmbeddingModel("OpenAIEmbedding", NAME_OPENAI_3SMALL, 1536, OpenAIEmbeddingImpl.INSTANCE, 500, 16);

    /**
     * 基于jbolt方案在本地部署的bge embedding模型，适用中文，1024 维。
     */
    public static final EmbeddingModel LOCAL_BGE_LARGE_ZH = new EmbeddingModel("local-bge-large-zh", NAME_BGE_LARGE_ZH_LOCAL, 1024, LocalEmbeddingImpl.INSTANCE, 500, 16);

    /**
     * 基于Ollama在本地部署的bge embedding模型，适用中文，1024 维。
     */
//    public static final EmbeddingModel OLLAMA_BGE_LARGE_ZH = new EmbeddingModel("ollama-bge-large-zh", NAME_BGE_LARGE_ZH_OLLAMA, 1024, OllamaEmbeddingImpl.INSTANCE, 500, 1);



    /**
     * 基于百度文心大模型的Embedding模型，适用中文，384 维。
     */
    public static final EmbeddingModel WEN_XIN = new EmbeddingModel("百度文心Embedding", NAME_BAIDU_QIANFAN, 384, WenxinEmbeddingImpl.INSTANCE, 384, 16);

    /**
     * 需要apikey和secretkey。由智源研究院研发的中文版文本表示模型，使用的是百度千帆云服务提供的接口。https://cloud.baidu.com/doc/WENXINWORKSHOP/s/dllz04sro
     */
    public static final EmbeddingModel BGE_LARGE_ZH = new EmbeddingModel("bge-large-zh", NAME_BGE_LARGE_ZH, 1024, WenxinEmbeddingImpl.INSTANCE, 500, 16);
    /**
     * 需要apikey和secretkey。由智源研究院研发的英文版文本表示模型，使用的是百度千帆云服务提供的接口。https://cloud.baidu.com/doc/WENXINWORKSHOP/s/mllz05nzk
     */
    public static final EmbeddingModel BGE_LARGE_EN = new EmbeddingModel("bge-large-en", NAME_BGE_LARGE_EN, 1024, WenxinEmbeddingImpl.INSTANCE, 500, 16);


    /**
     * 腾讯向量数据库专用，适用中文，768 维，推荐使用。
     */
    public static final EmbeddingModel TENCENT_BGE_BASE_ZH = new EmbeddingModel("腾讯BGE_BASE_ZH", NAME_TENCENT_BGE_BASE_ZH, 768, null, 500, 1000);
    /**
     * 腾讯向量数据库专用，适用中文，768 维。
     */
    public static final EmbeddingModel TENCENT_M3E_BASE = new EmbeddingModel("腾讯M3E_BASE", NAME_TENCENT_M3E_BASE, 768, null, 500, 1000);
    /**
     * 腾讯向量数据库专用，适用中文，1024 维。
     */
    public static final EmbeddingModel TENCENT_E5_LARGE_V2 = new EmbeddingModel("腾讯BGE_E5_LARGE_V2", NAME_TENCENT_E5_LARGE_V2, 1024, null, 500, 1000);
    /**
     * 腾讯向量数据库专用,适用中文，1024 维。
     */
    public static final EmbeddingModel TENCENT_TEXT2VEC_LARGE_CHINESE = new EmbeddingModel("腾讯TEXT2VEC_LARGE_CHINESE", NAME_TENCENT_TEXT2VEC_LARGE_CHINESE, 1024, null, 500, 1000);
    /**
     * 腾讯向量数据库专用,适用于多种语言类型，768 维。
     */
    public static final EmbeddingModel TENCENT_MULTILINGUAL_E5_BASE = new EmbeddingModel("腾讯MULTILINGUAL_E5_BASE", NAME_TENCENT_MULTILINGUAL_E5_BASE, 768, null, 500, 1000);

    /**
     * 自定义的大模型
     */
    private static final ConcurrentHashMap<String, EmbeddingModel> CUSTOM_MODELS = new ConcurrentHashMap<>();

    /**
     * 注册自定义的Embedding大模型类型
     * @param model
     * @return
     */
    public static void registerCustomModel(EmbeddingModel model) {
        if (StrUtil.isBlank(model.getName())) {
            throw new IllegalArgumentException("自定义Embedding大模型缺少name属性，请检查");
        }
        EmbeddingModel embeddingModel = CUSTOM_MODELS.putIfAbsent(model.getName(), model);
        if (embeddingModel != null) {
            throw new RuntimeException("自定义Embedding大模型：" + model.getName() + "已注册，不可重复注册");
        }

    }

    /**
     * 根据名字获取注册的自定义Embedding大模型
     * @param name
     * @return
     */
    public static EmbeddingModel getCustomModel(String name) {
        return CUSTOM_MODELS.get(name);
    }

    /**
     * 移除自定义的Embedding大模型
     * @param name
     */
    public static void removeCustomModel(String name) {
        CUSTOM_MODELS.remove(name);
    }

    /**
     * 根据名字获取Embedding大模型
     * @param name
     * @return
     */
    public static EmbeddingModel getModel(String name) {
        switch (name) {
            case EmbeddingModel.NAME_BAI_CHUAN:
                return EmbeddingModel.BAI_CHUAN;
            case EmbeddingModel.NAME_OPENAI_ADA002:
                return EmbeddingModel.OPENAI_ADA002;
            case EmbeddingModel.NAME_OPENAI_3SMALL:
                return EmbeddingModel.OPENAI_3_SMALL;
            case EmbeddingModel.NAME_TONGYI_V2:
                return EmbeddingModel.TONGYI_V2;
            case EmbeddingModel.NAME_TONGYI_V3:
                return EmbeddingModel.TONGYI_V3;
            case EmbeddingModel.NAME_TENCENT_BGE_BASE_ZH:
                return EmbeddingModel.TENCENT_BGE_BASE_ZH;
            case EmbeddingModel.NAME_TENCENT_M3E_BASE:
                return EmbeddingModel.TENCENT_M3E_BASE;
            case EmbeddingModel.NAME_TENCENT_E5_LARGE_V2:
                return EmbeddingModel.TENCENT_E5_LARGE_V2;
            case EmbeddingModel.NAME_TENCENT_TEXT2VEC_LARGE_CHINESE:
                return EmbeddingModel.TENCENT_TEXT2VEC_LARGE_CHINESE;
            case EmbeddingModel.NAME_TENCENT_MULTILINGUAL_E5_BASE:
                return EmbeddingModel.TENCENT_MULTILINGUAL_E5_BASE;
            case EmbeddingModel.NAME_BGE_LARGE_ZH_LOCAL:
                return EmbeddingModel.LOCAL_BGE_LARGE_ZH;
            default:
                return getCustomModel(name);


        }
    }


    /**
     * 是否是腾讯内部模型
     *
     * @param model
     * @return
     */
    public static boolean isTencentInnerModel(EmbeddingModel model) {
        return model == TENCENT_BGE_BASE_ZH || model == TENCENT_M3E_BASE || model == TENCENT_E5_LARGE_V2 || model == TENCENT_TEXT2VEC_LARGE_CHINESE || model == TENCENT_MULTILINGUAL_E5_BASE;
    }


    public EmbeddingModel(String name, String defaultModel, int dimension, EmbeddingProcessor processor, int chunkTokenLimit, int chunkCountLimit) {
        this.name = name;
        this.defaultModel = defaultModel;
        this.dimension = dimension;
        this.processor = processor;
        this.chunkCountLimit = chunkCountLimit;
        this.chunkTokenLimit = chunkTokenLimit;
    }





    /**
     * 向量化后的维度
     */
    private int dimension;

    /**
     * 模型名称
     */
    private String name;

    /**
     * 默认使用的模型
     */
    private String defaultModel;

    /**
     * 向量化的每个文本块长度不超过的token数，一般一个汉字算1个token，一个单词算1.3。
     */
    private int chunkTokenLimit;

    /**
     * 向量化一次性传入的文本数量不超过多少个
     */
    private int chunkCountLimit;

    /**
     * 该模型的处理实现类
     */
    private EmbeddingProcessor processor;


    public int getDimension() {
        return dimension;
    }

    public String getName() {
        return name;
    }

    public String getDefaultModel() {
        return defaultModel;
    }

    public int getChunkTokenLimit() {
        return chunkTokenLimit;
    }

    public int getChunkCountLimit() {
        return chunkCountLimit;
    }

    public EmbeddingProcessor getProcessor() {
        return processor;
    }
}
